# -*- coding:utf-8 -*-

import torch
from tqdm import tqdm
from utils import read_data, MyDataset
from config import parsers
from torch.utils.data import DataLoader
from model import MyModel
from sklearn.metrics import accuracy_score, f1_score, recall_score, precision_score
import logging
import pandas as pd
from log import logger_init

label_map_1 = {'服务提醒': 0, "内容资讯": 1, "营销活动": 2, "社交通讯": 3}
label_map_2 = {"私信-社交通讯-即时聊天": 0, "私信-服务提醒-订阅": 1, "私信-服务提醒-出行": 2, "私信-服务提醒-健康": 3, "私信-服务提醒-工作事项提醒": 4,
               "私信-服务提醒-帐号动态": 5, "私信-服务提醒-订单和物流": 6, "私信-服务提醒-财务": 7, "私信-服务提醒-设备提醒": 8, "私信-服务提醒-系统提示": 9,
               "私信-服务提醒-邮件": 10, "私信-服务提醒-正在发生的事": 11, "公信-内容资讯-内容推荐": 12, "公信-内容资讯-新闻": 13, "公信-内容资讯-财经动态": 14,
               "公信-内容资讯-生活资讯": 15,
               "公信-内容资讯-社交动态": 16, "公信-内容资讯-调研": 17, "公信-内容资讯-其他": 18, "公信-营销活动-产品促销": 19, "公信-营销活动-功能推荐": 20,
               "公信-营销活动-运营活动": 21}


def output_test_file(tru, pre, test_text):
    label_map = {v: k for k, v in label_map_2.items()}
    diff_index = []
    for i in range(len(pre)):
        if pre[i] != tru[i]:
            # print(f"第{i}项不同，True为{tru[i]}, 模型给出的结果为{pre[i]}")
            diff_index.append(i)
    true_res = [label_map[tru[i]] for i in diff_index]
    pred_res = [label_map[pre[i]] for i in diff_index]
    diff_push = [test_text[i] for i in diff_index]
    df = pd.DataFrame({"真实标签": true_res, "模型推理标签": pred_res, "对应推文": diff_push})
    df.to_csv('logs/result_22.csv.txt')


def test_data():
    args = parsers()
    logger_init(log_level=logging.INFO)
    device = "cuda:0" if torch.cuda.is_available() else "cpu"

    test_text, test_label = read_data(args.test_file)
    test_dataset = MyDataset(test_text, test_label, args.max_len)
    test_dataloader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)

    model = MyModel().to(device)
    model.load_state_dict(torch.load(args.save_model_best))
    model.eval()

    all_pred, all_true = [], []
    with torch.no_grad():
        for batch_text, batch_label in tqdm(test_dataloader):
            batch_label, batch_label = batch_label.to(device), batch_label.to(device)
            pred = model(batch_text)
            pred = torch.argmax(pred, dim=1)

            pred = pred.cpu().numpy().tolist()
            label = batch_label.cpu().numpy().tolist()
            all_pred.extend(pred)
            all_true.extend(label)

    accuracy = accuracy_score(all_true, all_pred)
    precision = precision_score(all_true, all_pred, average="micro")
    recall = recall_score(all_true, all_pred, average="micro")
    f1 = f1_score(all_true, all_pred, average="micro")
    output_test_file(all_true, all_pred, test_text)
    logging.info(f"test dataset accuracy:{accuracy:.4f}\tprecision:{precision:.4f}\trecall:{recall:.4f}\tf1:{f1:.4f}")


if __name__ == "__main__":
    test_data()
